# Copyright (c) 2023 ETH Zurich.
#                    All rights reserved.
#
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# The source code is adapted from the sorting source code written by
# Nils Blach.
#
# main author: Robert Gerstenberger

from typing import Dict, List, Set


def string_to_list(string: str) -> List[int]:
    """
    Helper function to convert a list encoded inside a string into a Python
    list object of integer elements.

    :param string: Input string containing a list.
    :type string: str
    :return: List of integer elements.
    :rtype: List[int]
    :raise AssertionError: If input string does not contain a list.
    """

    assert string[0] == "[" and string[-1] == "]", "String is not a list."
    return [int(num) for num in string[1:-1].split(",")]


def string_to_set(string: str) -> Set[int]:
    """
    Helper function to convert a list encoded inside a string into a Python
    set object of integer elements.

    :param string: Input string containing a list.
    :type string: str
    :return: Set of integer elements.
    :rtype: Set[int]
    :raise AssertionError: If input string does not contain a list.
    """

    assert string[0] == "[" and string[-1] == "]", "String is not a list."
    return {int(num) for num in string[1:-1].split(",")}


def test_set_intersection(state: Dict) -> bool:
    """
    Function to test whether the final solution matches ground truth.

    :param state: Thought state that represents the final solution.
    :type state: Dict
    :return: Returns whether the solution matches the ground truth.
    :rtype: bool
    """

    # convert string to list
    try:
        correct_list = string_to_list(state["result"])
        sorted_list = sorted(string_to_list(state["current"]))
        return sorted_list == correct_list
    except:
        return False


def num_errors(state_or_states):
    """
    兼容单个Dict或List[Dict]，返回float或List[float]，用于Score操作。
    """
    def _num_errors(state: Dict) -> float:
        # print(f"=== 调试num_errors函数 ===")
        # print(f"输入state: {state}")
        try:
            set1 = string_to_set(state["set1"])
            set2 = string_to_set(state["set2"])
            # print(f"set1: {set1}")
            # print(f"set2: {set2}")
            
            if "subset" in state and state["subset"] != "" and state["subset"] is not None:
                subset_str = state["subset"]
                if "@" in subset_str:
                    subset_parts = subset_str.split("@")
                    set1 = string_to_set(subset_parts[0])
                    set2 = string_to_set(subset_parts[1])
                    # print(f"使用subset分割@，set1: {set1}, set2: {set2}")
                else:
                    set2 = string_to_set(subset_str)
                    # print(f"使用subset作为set2: {set2}")
            
            common = sorted(list(set1 & set2))
            # print(f"正确交集: {common}")
            
            llm_solution = sorted(string_to_list(state["current"]))
            # print(f"LLM解答: {llm_solution}")
            
            num_errors = 0
            common_idx = 0
            llm_idx = 0
            while common_idx < len(common) and llm_idx < len(llm_solution):
                if common[common_idx] == llm_solution[llm_idx]:
                    common_idx += 1
                    llm_idx += 1
                elif common[common_idx] < llm_solution[llm_idx]:
                    common_idx += 1
                    num_errors += 1
                elif common[common_idx] > llm_solution[llm_idx]:
                    llm_idx += 1
                    num_errors += 1
            num_errors += len(common) - common_idx + len(llm_solution) - llm_idx
            # print(f"计算出的错误数: {num_errors}")
            return num_errors
        except Exception as e:
            print(f"num_errors函数出现异常: {e}")
            print(f"异常类型: {type(e)}")
            import traceback
            print(f"异常堆栈: {traceback.format_exc()}")
            return 1000
    
    if isinstance(state_or_states, list):
        # print(f"处理多个states，数量: {len(state_or_states)}")
        return [_num_errors(s) for s in state_or_states]
    else:
        # print(f"处理单个state")
        return _num_errors(state_or_states)